from langchain_community.utilities import SQLDatabase
from langchain_ollama import ChatOllama
from src.config import (
    ollama_host,
    ollama_port,
    ollama_model,
    milvus_url,
    model_embedding,
    milvus_mysql_collection,
    milvus_limit,
    milvus_db_name,
    mysql_ip,
    mysql_user_name,
    mysql_password,
    mysql_port,
    mysql_db, milvus_tool_collection
)
import ollama
from pymilvus import MilvusClient


def create_chat(ollama_model="qwen2.5:7b", temperature=0):
    llm = ChatOllama(
        base_url=ollama_host + ':' + str(ollama_port),
        model=ollama_model,
        temperature=temperature,
        keep_alive=1
    )
    return llm


def create_chat_with_options(ollama_model="qwen2.5:7b", options={}, params={}):
    default_options = {
        "temperature": 0
    }
    options = {**default_options, **options}

    llm = ChatOllama(
        base_url=ollama_host + ':' + str(ollama_port),
        model=ollama_model,
        keep_alive=1,
        **options,
        ** params
    )
    return llm


client = MilvusClient(
    uri=milvus_url,
    db_name=milvus_db_name,
)


def embedding(query):
    llm_client = ollama.Client(host=ollama_host + ':' + str(ollama_port))
    v = llm_client.embeddings(
        model=model_embedding,
        prompt=query
    )
    return v.embedding


def query_v_from_milvus_by_vector(collection_name, vector, limit, output_fields):
    res = client.search(collection_name=collection_name,
                        data=[vector],
                        limit=limit,
                        output_fields=output_fields)
    return list(res)


def query_mysql_pk(query):
    res = query_v_from_milvus_by_vector(
        collection_name=milvus_mysql_collection,
        vector=embedding(query),
        limit=milvus_limit,
        output_fields=["id", "text", "model"]
    )
    return res


def query_tools(query):
    res = query_v_from_milvus_by_vector(
        collection_name=milvus_tool_collection,
        vector=embedding(query),
        limit=milvus_limit,
        output_fields=["desc", "tool_name"]
    )
    return res


def create_db():
    db = SQLDatabase.from_uri(f'mysql+pymysql://{mysql_user_name}:{mysql_password}@{mysql_ip}:{mysql_port}/{mysql_db}')
    return db


if __name__ == '__main__':
    querys = ["查询2025年3月15日和2025年3月17日三江化工和嘉兴华雯化工上游来水的pH和cod数据",

                "查询2025年3月15日到2025年3月17日三江化工和嘉兴华雯化工上游来水以及下游排海的ph,cod数据",
                "查询2025年3月10日到2025年3月16日上游来水的数据",
                "查询2025年3月10日到2025年3月16日下游排海的数据",
                "查询2025年3月9日到2025年3月12日厂内在线数据",
                "查询2025年3月9日到11日实验室检测二沉池(A)、二沉池(B)和磁混凝的异常数据",
                "查询厂内运行巡检打卡记录数据",
                "查询运行缺卡数据",
                "查询2025年3月10日到2025年3月16日的2#机主和深度处理A的用电情况",
                "查询2025年3月10日到2025年3月16日的生活用水量",
                "查询2025年1月1日到2025年3月16日的用蒸汽情况",
                "查询2025年3月10日到3月19日营养液、液氧的药剂出入库情况",
                "查询2025年3月10日到3月19日营养液、液氧的药剂总量情况",
                "查询药剂液氮入库情况",
                "入库时间为2025年3月污泥外运平湖石化情况",
                "查询2025年3月20日运行打卡数据",
                "查询2025年3月16日到23日运行缺卡数据",
                "查询考勤数据",
                "查询2025年3月考勤打卡数据",
                "获取三江化工和传化合成的今天上游来水和下游排海"]
    for query in querys:
        res = query_mysql_pk(query)
        print(query+ ": ", res[0][0]["entity"]["model"])
    # res = query_v_from_milvus_by_vector(
    #     collection_name="tool_map_desc",
    #     vector=embedding("查询2025年3月10日到3月19日营养液、液氧的药剂出入库情况"),
    #     limit=milvus_limit,
    #     output_fields=["tool_name"]
    # )
    # print([i.get("entity", None).get("tool_name", None) for i in res[0]])
